package pinger

import (
	"bytes"
	"fmt"
	"net"
	"os"
	"time"

	"gitee.com/sillyman/simpleUtil/common/empty"
	"gitee.com/sillyman/simpleUtil/network/ip4util"
	"github.com/pkg/errors"
	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
	"golang.org/x/net/ipv6"
)

// pingId 进程ID
var (
	pingId = os.Getpid() & 0xffff

	// ErrTimeout 超时
	ErrTimeout       = errors.New("i/o timeout")
	ErrUserInterrupt = errors.New("user interrupt")
)

// Pinger
type Pinger struct {
	// Timeout 超时时间
	Timeout time.Duration

	// PingInterval 下一次Ping的间隔
	PingInterval time.Duration

	// TryTimes 尝试次数
	TryTimes int

	// PayloadSize 荷载数据的长度
	PayloadSize int

	// OnRecv 收包的回调函数
	OnRecv func(r *ReplyInfo)

	// interruptCh 用户终止正在的 Pinging
	interruptCh chan empty.T
}

// New 创建一个Pinger实例
func New() *Pinger {
	return &Pinger{
		Timeout:      2 * time.Second,
		PingInterval: time.Second,
		TryTimes:     2,
		PayloadSize:  22,
		OnRecv: func(r *ReplyInfo) {
			fmt.Println(r.String())
		},
		interruptCh: make(chan empty.T),
	}
}

// Ping
// times Ping多少次
func (p *Pinger) Ping(target net.IP, times int) error {
	isIPv6 := false
	lProto, lAddr := "ip4:icmp", "0.0.0.0"
	if !ip4util.IsIPv4(target) {
		lProto, lAddr = "ip6:ipv6-icmp", "::"
		isIPv6 = true
	}

	conn, err := icmp.ListenPacket(lProto, lAddr)
	if err != nil {
		return err
	}
	defer conn.Close()

	for i := 0; i < times; i++ {
		select {
		case <-p.interruptCh:
			return nil

		case <-time.After(p.PingInterval): //间隔时间
			p.pingOne(conn, target, i, isIPv6)
		}
	}

	return nil
}

// pingOne ping一次目标主机
func (p *Pinger) pingOne(conn *icmp.PacketConn, target net.IP, seq int, isIPv6 bool) *ReplyInfo {
	var (
		timeOfRecv  time.Time
		timeOfSend  time.Time
		recvReplyCh = make(chan *ReplyInfo)
		r           = &ReplyInfo{From: target, IsIPv6: isIPv6, PayloadSize: p.PayloadSize, Seq: seq}
	)

	// 回调函数
	defer func() {
		if p.OnRecv != nil {
			p.OnRecv(r)
		}
	}()

	// reply
	go func(recvCh chan<- *ReplyInfo) {
		defer func() { // 关闭资源
			close(recvReplyCh)
		}()
		for i := 0; i < p.TryTimes; i++ {
			typ, code, err := p.recvReply(conn, target, seq, isIPv6)
			if err != nil {
				if errors.Is(ErrTimeout, err) && i < p.TryTimes-1 {
					continue
				}
				r.Err = err
				recvCh <- r
				return

			} else {
				timeOfRecv = time.Now()
				r.ICMPType, r.ICMPCode = typ, code
				recvCh <- r
				return
			}
		}
	}(recvReplyCh)

	// request
	for i := 0; i < p.TryTimes; i++ {
		r.Retry = i

		if i == 0 { // 第一需要立即发送
			if err := p.sendRequest(conn, target, seq, isIPv6); err != nil {
				r.Err = errors.Wrap(err, "failed to send icmp-echo")
				return r
			}
			timeOfSend = time.Now()
		}

		select {
		case r := <-recvReplyCh:
			r.RTT = timeOfRecv.Sub(timeOfSend)
			return r // 成功

		case <-time.After(p.PingInterval): // 之后的重试需要等待间隔时间
			if err := p.sendRequest(conn, target, seq, isIPv6); err != nil {
				r.Err = errors.Wrap(err, "failed to send icmp-echo")
				return r
			}
			timeOfSend = time.Now()
		}
	}

	r.Err = ErrTimeout
	return r
}

// sendRequest 发送ICMP请求
func (p *Pinger) sendRequest(conn *icmp.PacketConn, target net.IP, seq int, isIPv6 bool) error {
	var reqTyp icmp.Type = ipv4.ICMPTypeEcho
	if isIPv6 {
		reqTyp = ipv6.ICMPTypeEchoRequest
	}

	echoMsg := &icmp.Message{
		Type: reqTyp, Code: 0,
		Body: &icmp.Echo{
			ID: pingId, Seq: seq,
			Data: bytes.Repeat([]byte{0xff}, p.PayloadSize),
		},
	}

	echoMsgBin, err := echoMsg.Marshal(nil)
	if err != nil {
		return err
	}

	if err := conn.SetWriteDeadline(time.Now().Add(p.Timeout)); err != nil {
		return err
	}

	_, err = conn.WriteTo(echoMsgBin, &net.IPAddr{IP: target})
	return err
}

// recvReply 接受ICMP应答
// 返回 icmp-type, icmp-code, err
func (p *Pinger) recvReply(conn *icmp.PacketConn, from net.IP, seq int, isIPv6 bool) (int, int, error) {
	var (
		buf   = make([]byte, 160+p.PayloadSize)
		after = time.After(p.Timeout)
	)

	proto := ipv4.ICMPTypeEcho.Protocol()
	if isIPv6 {
		proto = ipv6.ICMPTypeEchoRequest.Protocol()
	}

For:
	for {
		select {
		case <-p.interruptCh:
			return 0, 0, ErrUserInterrupt
		case <-after:
			return 0, 0, ErrTimeout
		default:
			if err := conn.SetReadDeadline(time.Now().Add(p.Timeout)); err != nil {
				return 0, 0, err
			}

			n, peer, err := conn.ReadFrom(buf)
			if err != nil {
				return 0, 0, err
			}

			if addr, ok := peer.(*net.IPAddr); !ok || !from.Equal(addr.IP) {
				continue For // 比较 IP 地址
			}

			replyMsg, err := icmp.ParseMessage(proto, buf[:n])
			if err != nil {
				continue For // 错误的报文，不能被解析
			}

			if replyMsg.Type == ipv4.ICMPTypeEcho || replyMsg.Type == ipv6.ICMPTypeEchoRequest {
				continue For // 如果是ping自己，则先收到 echo，需要跳过
			}

			replyMsgBody, ok := replyMsg.Body.(*icmp.Echo)
			if !ok || replyMsgBody == nil {
				continue For // 不完整的报文
			}

			if replyMsgBody.ID != pingId {
				continue For // 比较id
			}

			if replyMsgBody.Seq != seq {
				continue For // 比较seq
			}

			var icmpTyp int
			if isIPv6 {
				icmpTyp = int(replyMsg.Type.(ipv6.ICMPType))
			} else {
				icmpTyp = int(replyMsg.Type.(ipv4.ICMPType))
			}
			return icmpTyp, replyMsg.Code, nil
		}
	}
}

// Interrupt 打断正在进行的 Pinging
func (p *Pinger) Interrupt() {
	select {
	case <-p.interruptCh:
	default:
		close(p.interruptCh)
	}
}
